#!/usr/bin/env python3
#coding:utf-8

__author__ = 'xmxoxo<xmxoxo@qq.com>'

# 识别图像中的人脸，并截取保存到文件

import argparse
import os
import sys
import numpy as np
import dlib
import cv2
import pandas as pd

from faces_lib import *
from tqdm import tqdm

def main():
    print('正在加载模型...')
    saveface = save
    # 加载模型
    detector, predictor = load_model()

    while 1:
        pl()
        filename = input('请输入图像文件名(Q退出):')
        filename = filename.strip()
        ## exit
        if filename in ['quit','q','Q']:
            break
        if not filename:continue
        if not os.path.exists(filename): continue


        # 检测图像,
        face_detector(filename, detector, predictor,
                    showwin=1, saveface=saveface,
                    pre_name='image_cap', outpath='./')

        cv2.destroyAllWindows()

def batch_pick(imagefile, outpath):
    '''按文件清单批量提取头像
    '''
    # 如果是文件，表示是图像文件清单
    if os.path.isfile(imagefile):
        dats = readtxt(imagefile).splitlines()

    # 如果是目录，则生成文件
    if os.path.isdir(imagefile):
        dats = [x[0] for x in get_files(imagefile, folder=0)]

    if not dats:
        print('未找到文件清单,退出...')
        return
    print('共有%d个文件待处理...'% len(dats))
    #print(dats[:10])
    #return
    rets = []
    # 加载模型
    detector, predictor = load_model()

    for fname in tqdm(dats, ncols=70):
        opath, fn = os.path.split(fname)
        pre_name = os.path.splitext(fn)[0]
        # 生成子目录
        npath = os.path.join(outpath, os.path.basename(opath))
        mkfold(npath)

        # 检测图像并保存
        ret = face_detector(fname, detector, predictor,
                    showwin=0, saveface=1,
                    pre_name=pre_name, outpath=npath)
        # 记录文件名以及 提取到人脸的个数
        rets.append([fname, len(ret)])
        if len(rets)>=100:
            pass

    # 保存记录的结果
    outfile = os.path.join(outpath, 'allfiles.csv')
    df_out = pd.DataFrame(rets, columns=["filename", "hasface"])
    df_out.to_csv(outfile)


def server():
    import streamlit as st

    st.set_page_config(
        page_title="人脸识别管理端",
        page_icon="favicon.png",
        layout="wide",
        menu_items={
            "About": "关于人脸识别系统",
        },
    )

    # 加载模型
    detector, predictor = load_model()

    st.header('人脸识别')
    st.text('功能说明: 人脸识别')
    st.write('---')

    uploaded_file = st.file_uploader("上传图像:", type=['jpg','png','jpeg'])
    if uploaded_file is None:
        st.write('上传后可预览。')
    else:

        st.write('---')
        scol1, scol2 = st.columns(2)
        with scol1:
            # 显示已选的文件
            st.write('上传预览:')
            st.image(uploaded_file, caption='已上传文件', width=400) #, use_column_width=True
            img_bin = uploaded_file.read()
            print('uploaded_file:', uploaded_file)
            print('image:', type(img_bin))

        with scol2:

            # 检测图像,
            faces, image = face_detector(img_bin, detector, predictor,
                        showwin=0, saveface=0,
                        pre_name='image_cap', outpath='./')
            st.write('识别结果:')
            st.image(image, caption='识别结果图像', width=400) #, use_column_width=True

        st.write('---')
        st.write("人脸提取结果:")
        print('faces:', type(faces))
        for face in faces:
            print('face:', type(face))
            st.image(face)#, width=200)

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='face_detect')
    parser.add_argument('--image', type=str, default="", help='单个图像文件')
    parser.add_argument('--imagefile', type=str, default="", help='图片文件清单')
    parser.add_argument('--outpath', type=str, default="", help='提取结果保存的目录')
    parser.add_argument('--save', type=int, default=0, help='是否保存图像')
    args = parser.parse_args()
    image = args.image
    imagefile = args.imagefile
    outpath = args.outpath
    save = args.save

    filename = ''
    if filename:
        #filename = imagefile
        #print(os.path.exists(filename))

        # 加载模型
        detector, predictor = load_model()
        # 检测图像
        face_detector(filename, detector, predictor)

        # 调用保存
        #face_detector(filename, detector, predictor, showwin=0, pre_name='image_out')
    else:
        # 判断是否批量调用
        if imagefile != '' and outpath != '':
            batch_pick(imagefile, outpath)
        else:
            #　进入交互模式
            # main()
            server()

# python face_find.py --image=./test/pp.jpg


